from functools import partial, update_wrapper
# ! <<<
# from keras.layers import *
# from keras.callbacks import *
# from keras.models import *
# from keras.optimizers import *
# from keras import backend as K
from tensorflow.keras.layers import *
from tensorflow.keras.callbacks import *
from tensorflow.keras.models import *
from tensorflow.keras.optimizers import *
from tensorflow.keras import backend as K
# ! >>>

# ! <<<
# import warnings, keras, sys
import warnings, sys
from tensorflow import keras
# ! >>>
import numpy as np
sys.dont_write_bytecode = True
"""-------------------------------------------------------2D Unet------------------------------------------------"""
def get_2D_Unet(inputshape,nclass,activation='sigmoid',axis=1):
    inputdata = Input(shape=inputshape)
    print ("Input data shape:",inputdata.shape)

    conv1_1 = Conv2D(32, 3, activation = 'relu', padding = 'same', kernel_initializer = 'he_normal', name = 'conv1_1')(inputdata)
    conv1_1_bn = BatchNormalization(name = 'conv1_1_bn')(conv1_1)
    conv1_2 = Conv2D(32, 3, activation = 'relu', padding = 'same', kernel_initializer = 'he_normal', name = 'conv1_2')(conv1_1_bn)
    conv1_2_bn = BatchNormalization(name = 'conv1_2_bn')(conv1_2)
    pool1 = MaxPooling2D(pool_size=(2, 2), name = 'pool1')(conv1_2_bn)
    print ("pool1 shape:",pool1.shape)

    conv2_1 = Conv2D(64, 3, activation = 'relu', padding = 'same', kernel_initializer = 'he_normal', name = 'conv2_1')(pool1)
    conv2_1_bn = BatchNormalization(name = 'conv2_1_bn')(conv2_1)
    conv2_2 = Conv2D(64, 3, activation = 'relu', padding = 'same', kernel_initializer = 'he_normal', name = 'conv2_2')(conv2_1_bn)
    conv2_2_bn = BatchNormalization(name = 'conv2_2_bn')(conv2_2)
    pool2 = MaxPooling2D(pool_size=(2, 2), name = 'pool2')(conv2_2_bn)
    print ("pool2 shape:",pool2.shape)

    conv3_1 = Conv2D(96, 3, activation = 'relu', padding = 'same', kernel_initializer = 'he_normal', name = 'conv3_1')(pool2)
    conv3_1_bn = BatchNormalization(name = 'conv3_1_bn')(conv3_1)
    conv3_2 = Conv2D(96, 3, activation = 'relu', padding = 'same', kernel_initializer = 'he_normal', name = 'conv3_2')(conv3_1_bn)
    conv3_2_bn = BatchNormalization(name = 'conv3_2_bn')(conv3_2)
    pool3 = MaxPooling2D(pool_size=(2, 2), name = 'pool3')(conv3_2_bn)
    print ("pool3 shape:",pool3.shape)

    conv4_1 = Conv2D(128, 3, activation = 'relu', padding = 'same', kernel_initializer = 'he_normal', name = 'conv4_1')(pool3)
    conv4_1_bn = BatchNormalization(name = 'conv4_1_bn')(conv4_1)
    conv4_2 = Conv2D(128, 3, activation = 'relu', padding = 'same', kernel_initializer = 'he_normal', name = 'conv4_2')(conv4_1_bn)
    conv4_2_bn = BatchNormalization(name = 'conv4_2_bn')(conv4_2)
    pool4 = MaxPooling2D(pool_size=(2, 2),name = 'pool4')(conv4_2_bn)
    print ("pool4 shape:",pool4.shape)

    conv5_1 = Conv2D(256, 3, activation = 'relu', padding = 'same', kernel_initializer = 'he_normal', name = 'conv5_1')(pool4)
    conv5_1_bn = BatchNormalization(name = 'conv5_1_bn')(conv5_1)
    conv5_2 = Conv2D(256, 3, activation = 'relu', padding = 'same', kernel_initializer = 'he_normal', name = 'conv5_2')(conv5_1_bn)
    conv5_2_bn = BatchNormalization(name = 'conv5_2_bn')(conv5_2)
    print ("conv5 shape:",conv5_2_bn.shape)

    up6 = Conv2D(128, 2, activation = 'relu', padding = 'same', kernel_initializer = 'he_normal', name = 'up6')(UpSampling2D(size = (2,2),name = 'up6_c')(conv5_2_bn))
    up6_bn = BatchNormalization(name = 'up6_bn')(up6)
    merge6 = concatenate([conv4_2_bn,up6_bn], axis = axis, name = 'merge6')
    drop6 = Dropout(0.5, name='drop6')(merge6)
    conv6_1 = Conv2D(128, 3, activation = 'relu', padding = 'same', kernel_initializer = 'he_normal', name = 'conv6_1')(drop6)
    conv6_1_bn = BatchNormalization(name = 'conv6_1_bn')(conv6_1)
    conv6_2 = Conv2D(128, 3, activation = 'relu', padding = 'same', kernel_initializer = 'he_normal', name = 'conv6_2')(conv6_1_bn)
    conv6_2_bn = BatchNormalization(name = 'conv6_2_bn')(conv6_2)
    print ("conv6 shape:",conv6_2_bn.shape)

    up7 = Conv2D(96, 2, activation = 'relu', padding = 'same', kernel_initializer = 'he_normal', name = 'up7')(UpSampling2D(size = (2,2), name = 'up7_c')(conv6_2_bn))
    up7_bn = BatchNormalization(name = 'up7_bn')(up7)
    merge7 = concatenate([conv3_2_bn,up7_bn], axis = axis, name = 'merge7')
    drop7 = Dropout(0.5, name = 'drop7')(merge7)
    conv7_1 = Conv2D(96, 3, activation = 'relu', padding = 'same', kernel_initializer = 'he_normal', name = 'conv7_1')(drop7)
    conv7_1_bn = BatchNormalization(name = 'conv7_1_bn')(conv7_1)
    conv7_2 = Conv2D(96, 3, activation = 'relu', padding = 'same', kernel_initializer = 'he_normal', name = 'conv7_2')(conv7_1_bn)
    conv7_2_bn = BatchNormalization(name = 'conv7_2_bn')(conv7_2)
    print ("conv7 shape:",conv7_2_bn.shape)

    up8 = Conv2D(64, 2, activation = 'relu', padding = 'same', kernel_initializer = 'he_normal', name = 'up8')(UpSampling2D(size = (2,2), name = 'up8_c')(conv7_2_bn))
    up8_bn = BatchNormalization(name = 'up8_bn')(up8)
    merge8 = concatenate([conv2_2_bn,up8_bn], axis = axis, name = 'merge8')
    drop8 = Dropout(0.5, name = 'drop8')(merge8)
    conv8_1 = Conv2D(64, 3, activation = 'relu', padding = 'same', kernel_initializer = 'he_normal', name = 'conv8_1')(drop8)
    conv8_1_bn = BatchNormalization(name = 'conv8_1_bn')(conv8_1)
    conv8_2 = Conv2D(64, 3, activation = 'relu', padding = 'same', kernel_initializer = 'he_normal', name = 'conv8_2')(conv8_1_bn)
    conv8_2_bn = BatchNormalization(name = 'conv8_2_bn')(conv8_2)
    print ("conv8 shape:",conv8_2_bn.shape)

    up9 = Conv2D(32, 2, activation = 'relu', padding = 'same', kernel_initializer = 'he_normal', name = 'up9')(UpSampling2D(size = (2,2), name = 'up9_c')(conv8_2_bn))
    up9_bn = BatchNormalization(name = 'up9_bn')(up9)
    merge9 = concatenate([conv1_2_bn,up9_bn], axis = axis, name = 'merge9')
    drop9 = Dropout(0.5, name = 'drop9')(merge9)
    conv9_1 = Conv2D(32, 3, activation = 'relu', padding = 'same', kernel_initializer = 'he_normal', name = 'conv9_1')(drop9)
    conv9_1_bn = BatchNormalization(name = 'conv9_1_bn')(conv9_1)
    conv9_2 = Conv2D(32, 3, activation = 'relu', padding = 'same', kernel_initializer = 'he_normal', name = 'conv9_2')(conv9_1_bn)
    conv9_2_bn = BatchNormalization(name = 'conv9_2_bn')(conv9_2)
    print ("conv9 shape:",conv9_2_bn.shape)
    conv9_3 = Conv2D(nclass, 3, activation = 'relu', padding = 'same', kernel_initializer = 'he_normal', name = 'conv9_3')(conv9_2_bn)

    if activation == 'sigmoid':
        score = Conv2D(1, 1, activation = 'sigmoid', padding='same', kernel_initializer='he_normal',name = 'score')(conv9_3)
    elif activation == 'softmax':
        score = Conv2D(nclass, 1, activation = 'softmax', padding='same', kernel_initializer='he_normal',name = 'score')(conv9_3)

    print ("final output shape:",score.shape)
    model = Model(inputs = inputdata, outputs = score)
    return model

"""-----------------------------------------------3D Unet Models---------------------------------------------"""
def get_3D_Unet(inputshape,nclass,activation='sigmoid',axis=1):
    inputdata = Input(shape=inputshape)
    print ("Input data shape:", inputdata.shape)

    conv1_1 = Conv3D(32, 3, activation = 'relu', padding = 'same', kernel_initializer = 'he_normal', name = 'conv1_1')(inputdata)
    conv1_1_bn = BatchNormalization(name = 'conv1_1_bn')(conv1_1)
    conv1_2 = Conv3D(32, 3, activation = 'relu', padding = 'same', kernel_initializer = 'he_normal', name = 'conv1_2')(conv1_1_bn)
    conv1_2_bn = BatchNormalization(name = 'conv1_2_bn')(conv1_2)
    pool1 = MaxPooling3D(pool_size=(2, 2, 2), name = 'pool1')(conv1_2_bn)
    print ("pool1 shape:",pool1.shape)

    conv2_1 = Conv3D(64, 3, activation = 'relu', padding = 'same', kernel_initializer = 'he_normal', name = 'conv2_1')(pool1)
    conv2_1_bn = BatchNormalization(name = 'conv2_1_bn')(conv2_1)
    conv2_2 = Conv3D(64, 3, activation = 'relu', padding = 'same', kernel_initializer = 'he_normal', name = 'conv2_2')(conv2_1_bn)
    conv2_2_bn = BatchNormalization(name = 'conv2_2_bn')(conv2_2)
    pool2 = MaxPooling3D(pool_size=(2, 2, 2), name = 'pool2')(conv2_2_bn)
    print ("pool2 shape:",pool2.shape)

    conv3_1 = Conv3D(96, 3, activation = 'relu', padding = 'same', kernel_initializer = 'he_normal', name = 'conv3_1')(pool2)
    conv3_1_bn = BatchNormalization(name = 'conv3_1_bn')(conv3_1)
    conv3_2 = Conv3D(96, 3, activation = 'relu', padding = 'same', kernel_initializer = 'he_normal', name = 'conv3_2')(conv3_1_bn)
    conv3_2_bn = BatchNormalization(name = 'conv3_2_bn')(conv3_2)
    pool3 = MaxPooling3D(pool_size=(2, 2, 2), name = 'pool3')(conv3_2_bn)
    print ("pool3 shape:",pool3.shape)

    conv4_1 = Conv3D(128, 3, activation = 'relu', padding = 'same', kernel_initializer = 'he_normal', name = 'conv4_1')(pool3)
    conv4_1_bn = BatchNormalization(name = 'conv4_1_bn')(conv4_1)
    conv4_2 = Conv3D(128, 3, activation = 'relu', padding = 'same', kernel_initializer = 'he_normal', name = 'conv4_2')(conv4_1_bn)
    conv4_2_bn = BatchNormalization(name = 'conv4_2_bn')(conv4_2)
    pool4 = MaxPooling3D(pool_size=(2, 2, 2),name = 'pool4')(conv4_2_bn)
    print ("pool4 shape:",pool4.shape)

    conv5_1 = Conv3D(256, 3, activation = 'relu', padding = 'same', kernel_initializer = 'he_normal', name = 'conv5_1')(pool4)
    conv5_1_bn = BatchNormalization(name = 'conv5_1_bn')(conv5_1)
    conv5_2 = Conv3D(256, 3, activation = 'relu', padding = 'same', kernel_initializer = 'he_normal', name = 'conv5_2')(conv5_1_bn)
    conv5_2_bn = BatchNormalization(name = 'conv5_2_bn')(conv5_2)
    print ("conv5 shape:",conv5_2_bn.shape)

    up6 = Conv3D(128, 2, activation = 'relu', padding = 'same', kernel_initializer = 'he_normal', name = 'up6')(UpSampling3D(size = (2,2,2),name = 'up6_c')(conv5_2_bn))
    # up6 = Conv3D(128, 2, activation = 'relu', padding = 'same', kernel_initializer = 'he_normal', name = 'up6')(conv5_2_bn)
    up6_bn = BatchNormalization(name = 'up6_bn')(up6)
    merge6 = concatenate([conv4_2_bn,up6_bn], axis = axis, name = 'merge6')
    drop6 = Dropout(0.5, name='drop6')(merge6)
    conv6_1 = Conv3D(128, 3, activation = 'relu', padding = 'same', kernel_initializer = 'he_normal', name = 'conv6_1')(drop6)
    conv6_1_bn = BatchNormalization(name = 'conv6_1_bn')(conv6_1)
    conv6_2 = Conv3D(128, 3, activation = 'relu', padding = 'same', kernel_initializer = 'he_normal', name = 'conv6_2')(conv6_1_bn)
    conv6_2_bn = BatchNormalization(name = 'conv6_2_bn')(conv6_2)
    print ("conv6 shape:",conv6_2_bn.shape)

    up7 = Conv3D(96, 2, activation = 'relu', padding = 'same', kernel_initializer = 'he_normal', name = 'up7')(UpSampling3D(size = (2,2,2), name = 'up7_c')(conv6_2_bn))
    up7_bn = BatchNormalization(name = 'up7_bn')(up7)
    merge7 = concatenate([conv3_2_bn,up7_bn], axis = axis, name = 'merge7')
    drop7 = Dropout(0.5, name = 'drop7')(merge7)
    conv7_1 = Conv3D(96, 3, activation = 'relu', padding = 'same', kernel_initializer = 'he_normal', name = 'conv7_1')(drop7)
    conv7_1_bn = BatchNormalization(name = 'conv7_1_bn')(conv7_1)
    conv7_2 = Conv3D(96, 3, activation = 'relu', padding = 'same', kernel_initializer = 'he_normal', name = 'conv7_2')(conv7_1_bn)
    conv7_2_bn = BatchNormalization(name = 'conv7_2_bn')(conv7_2)
    print ("conv7 shape:",conv7_2_bn.shape)

    up8 = Conv3D(64, 2, activation = 'relu', padding = 'same', kernel_initializer = 'he_normal', name = 'up8')(UpSampling3D(size = (2,2,2), name = 'up8_c')(conv7_2_bn))
    up8_bn = BatchNormalization(name = 'up8_bn')(up8)
    merge8 = concatenate([conv2_2_bn,up8_bn], axis = axis, name = 'merge8')
    drop8 = Dropout(0.5, name = 'drop8')(merge8)
    conv8_1 = Conv3D(64, 3, activation = 'relu', padding = 'same', kernel_initializer = 'he_normal', name = 'conv8_1')(drop8)
    conv8_1_bn = BatchNormalization(name = 'conv8_1_bn')(conv8_1)
    conv8_2 = Conv3D(64, 3, activation = 'relu', padding = 'same', kernel_initializer = 'he_normal', name = 'conv8_2')(conv8_1_bn)
    conv8_2_bn = BatchNormalization(name = 'conv8_2_bn')(conv8_2)
    print ("conv8 shape:",conv8_2_bn.shape)

    up9 = Conv3D(32, 2, activation = 'relu', padding = 'same', kernel_initializer = 'he_normal', name = 'up9')(UpSampling3D(size = (2,2,2), name = 'up9_c')(conv8_2_bn))
    up9_bn = BatchNormalization(name = 'up9_bn')(up9)
    merge9 = concatenate([conv1_2_bn,up9_bn], axis = axis, name = 'merge9')
    drop9 = Dropout(0.5, name = 'drop9')(merge9)
    conv9_1 = Conv3D(32, 3, activation = 'relu', padding = 'same', kernel_initializer = 'he_normal', name = 'conv9_1')(drop9)
    conv9_1_bn = BatchNormalization(name = 'conv9_1_bn')(conv9_1)
    conv9_2 = Conv3D(32, 3, activation = 'relu', padding = 'same', kernel_initializer = 'he_normal', name = 'conv9_2')(conv9_1_bn)
    conv9_2_bn = BatchNormalization(name = 'conv9_2_bn')(conv9_2)
    print ("conv9 shape:",conv9_2_bn.shape)

    conv9_3 = Conv3D(nclass, 3, activation = 'relu', padding = 'same', kernel_initializer = 'he_normal', name = 'conv9_3')(conv9_2_bn)

    if activation == 'sigmoid':
        score = Conv3D(1, 1, activation = 'sigmoid', padding='same', kernel_initializer='he_normal',name = 'score')(conv9_3)
    elif activation == 'softmax':
        score = Conv3D(nclass, 1, activation = 'softmax', padding='same', kernel_initializer='he_normal',name = 'score')(conv9_3)

    print("final output shape:",score.shape)
    model = Model(inputs=inputdata, outputs=score)
    return model

def get_3D_ResUnet(inputshape, nclass, activation='sigmoid',axis=1):
    inputdata = Input(shape=inputshape)
    print("Input data shape:", inputdata.shape)

    conv1_1 = Conv3D(32, 3, activation='relu', padding='same', kernel_initializer='he_normal', name='conv1_1')(
        inputdata)
    conv1_1_bn = BatchNormalization(name='conv1_1_bn')(conv1_1)
    conv1_2 = Conv3D(32, 3, activation='relu', padding='same', kernel_initializer='he_normal', name='conv1_2')(
        conv1_1_bn)
    conv1_2_bn = BatchNormalization(name='conv1_2_bn')(conv1_2)

    convc1 = concatenate([inputdata, conv1_2_bn], axis=axis)
    pool1 = MaxPooling3D(pool_size=(2, 2, 2), name='pool1')(convc1)
    print("pool1 shape:", pool1.shape)

    conv2_1 = Conv3D(64, 3, activation='relu', padding='same', kernel_initializer='he_normal', name='conv2_1')(pool1)
    conv2_1_bn = BatchNormalization(name='conv2_1_bn')(conv2_1)
    conv2_2 = Conv3D(64, 3, activation='relu', padding='same', kernel_initializer='he_normal', name='conv2_2')(
        conv2_1_bn)
    conv2_2_bn = BatchNormalization(name='conv2_2_bn')(conv2_2)

    convc2 = concatenate([pool1, conv2_2_bn], axis=axis)
    pool2 = MaxPooling3D(pool_size=(2, 2, 2), name='pool2')(convc2)
    print("pool2 shape:", pool2.shape)

    conv3_1 = Conv3D(96, 3, activation='relu', padding='same', kernel_initializer='he_normal', name='conv3_1')(pool2)
    conv3_1_bn = BatchNormalization(name='conv3_1_bn')(conv3_1)
    conv3_2 = Conv3D(96, 3, activation='relu', padding='same', kernel_initializer='he_normal', name='conv3_2')(
        conv3_1_bn)
    conv3_2_bn = BatchNormalization(name='conv3_2_bn')(conv3_2)

    convc3 = concatenate([pool2, conv3_2_bn], axis=axis)
    pool3 = MaxPooling3D(pool_size=(2, 2, 2), name='pool3')(convc3)
    print("pool3 shape:", pool3.shape)

    conv4_1 = Conv3D(128, 3, activation='relu', padding='same', kernel_initializer='he_normal', name='conv4_1')(pool3)
    conv4_1_bn = BatchNormalization(name='conv4_1_bn')(conv4_1)
    conv4_2 = Conv3D(128, 3, activation='relu', padding='same', kernel_initializer='he_normal', name='conv4_2')(
        conv4_1_bn)
    conv4_2_bn = BatchNormalization(name='conv4_2_bn')(conv4_2)

    convc4 = concatenate([pool3, conv4_2_bn], axis=axis)
    pool4 = MaxPooling3D(pool_size=(2, 2, 2), name='pool4')(convc4)
    print("pool4 shape:", pool4.shape)

    conv5_1 = Conv3D(256, 3, activation='relu', padding='same', kernel_initializer='he_normal', name='conv5_1')(pool4)
    conv5_1_bn = BatchNormalization(name='conv5_1_bn')(conv5_1)
    conv5_2 = Conv3D(256, 3, activation='relu', padding='same', kernel_initializer='he_normal', name='conv5_2')(
        conv5_1_bn)
    conv5_2_bn = BatchNormalization(name='conv5_2_bn')(conv5_2)

    convc5 = concatenate([pool4, conv5_2_bn], axis=axis)
    print("convc5 shape:", convc5.shape)

    up6 = Conv3D(128, 2, activation='relu', padding='same', kernel_initializer='he_normal', name='up6')(
        UpSampling3D(size=(2, 2, 2), name='up6_c')(convc5))
    up6_bn = BatchNormalization(name='up6_bn')(up6)
    merge6 = concatenate([conv4_2_bn, up6_bn], axis=axis, name='merge6')
    drop6 = Dropout(0.5, name='drop6')(merge6)

    conv6_1 = Conv3D(128, 3, activation='relu', padding='same', kernel_initializer='he_normal', name='conv6_1')(drop6)
    conv6_1_bn = BatchNormalization(name='conv6_1_bn')(conv6_1)
    conv6_2 = Conv3D(128, 3, activation='relu', padding='same', kernel_initializer='he_normal', name='conv6_2')(
        conv6_1_bn)
    conv6_2_bn = BatchNormalization(name='conv6_2_bn')(conv6_2)

    convc6 = concatenate([drop6, conv6_2_bn], axis=axis)
    print("conv6 shape:", convc6.shape)

    up7 = Conv3D(96, 2, activation='relu', padding='same', kernel_initializer='he_normal', name='up7')(
        UpSampling3D(size=(2, 2, 2), name='up7_c')(convc6))
    up7_bn = BatchNormalization(name='up7_bn')(up7)
    merge7 = concatenate([conv3_2_bn, up7_bn], axis=axis, name='merge7')
    drop7 = Dropout(0.5, name='drop7')(merge7)

    conv7_1 = Conv3D(96, 3, activation='relu', padding='same', kernel_initializer='he_normal', name='conv7_1')(drop7)
    conv7_1_bn = BatchNormalization(name='conv7_1_bn')(conv7_1)
    conv7_2 = Conv3D(96, 3, activation='relu', padding='same', kernel_initializer='he_normal', name='conv7_2')(
        conv7_1_bn)
    conv7_2_bn = BatchNormalization(name='conv7_2_bn')(conv7_2)

    convc7 = concatenate([drop7, conv7_2_bn], axis=axis)
    print("conv7 shape:", convc7.shape)

    up8 = Conv3D(64, 2, activation='relu', padding='same', kernel_initializer='he_normal', name='up8')(
        UpSampling3D(size=(2, 2, 2), name='up8_c')(convc7))
    up8_bn = BatchNormalization(name='up8_bn')(up8)
    merge8 = concatenate([conv2_2_bn, up8_bn], axis=axis, name='merge8')
    drop8 = Dropout(0.5, name='drop8')(merge8)

    conv8_1 = Conv3D(64, 3, activation='relu', padding='same', kernel_initializer='he_normal', name='conv8_1')(drop8)
    conv8_1_bn = BatchNormalization(name='conv8_1_bn')(conv8_1)
    conv8_2 = Conv3D(64, 3, activation='relu', padding='same', kernel_initializer='he_normal', name='conv8_2')(
        conv8_1_bn)
    conv8_2_bn = BatchNormalization(name='conv8_2_bn')(conv8_2)

    convc8 = concatenate([drop8, conv8_2_bn], axis=axis)
    print("conv8 shape:", convc8.shape)

    up9 = Conv3D(32, 2, activation='relu', padding='same', kernel_initializer='he_normal', name='up9')(
        UpSampling3D(size=(2, 2, 2), name='up9_c')(convc8))
    up9_bn = BatchNormalization(name='up9_bn')(up9)
    merge9 = concatenate([conv1_2_bn, up9_bn], axis=axis, name='merge9')
    drop9 = Dropout(0.5, name='drop9')(merge9)

    conv9_1 = Conv3D(32, 3, activation='relu', padding='same', kernel_initializer='he_normal', name='conv9_1')(drop9)
    conv9_1_bn = BatchNormalization(name='conv9_1_bn')(conv9_1)
    conv9_2 = Conv3D(32, 3, activation='relu', padding='same', kernel_initializer='he_normal', name='conv9_2')(
        conv9_1_bn)
    conv9_2_bn = BatchNormalization(name='conv9_2_bn')(conv9_2)

    convc9 = concatenate([drop9, conv9_2_bn], axis=axis)
    print("conv9 shape:", convc9.shape)

    conv9_3 = Conv3D(nclass, 3, activation='relu', padding='same', kernel_initializer='he_normal', name='conv9_3')(
        convc9)

    if activation == 'sigmoid':
        score = Conv3D(1, 1, activation = 'sigmoid', padding='same', kernel_initializer='he_normal',name = 'score')(conv9_3)
    elif activation == 'softmax':
        score = Conv3D(nclass, 1, activation = 'softmax', padding='same', kernel_initializer='he_normal',name = 'score')(conv9_3)

    print("final output shape:", score.shape)
    model = Model(inputs=inputdata, outputs=score)
    return model
